# utils/visualization.py
"""
可视化工具模块
"""

import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class TradingVisualizer:
    """交易可视化类"""
    
    def __init__(self, save_dir: str = "charts"):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(exist_ok=True)
        
        # 设置样式
        sns.set_style("whitegrid")
        self.colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']
    
    def plot_price_and_indicators(self, data: pd.DataFrame, title: str = "价格和技术指标",
                                 save: bool = True) -> None:
        """绘制价格和技术指标"""
        fig, axes = plt.subplots(4, 1, figsize=(15, 16))
        
        # 1. 价格和移动平均线
        ax1 = axes[0]
        ax1.plot(data.index, data['close'], label='收盘价', linewidth=1.5, color=self.colors[0])
        
        if 'sma_5' in data.columns:
            ax1.plot(data.index, data['sma_5'], label='SMA(5)', alpha=0.7, color=self.colors[1])
        if 'sma_20' in data.columns:
            ax1.plot(data.index, data['sma_20'], label='SMA(20)', alpha=0.7, color=self.colors[2])
        
        ax1.set_title(f'{title} - 价格和移动平均线')
        ax1.set_ylabel('价格')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 2. 成交量
        ax2 = axes[1]
        ax2.bar(data.index, data['volume'], alpha=0.6, color=self.colors[3])
        ax2.set_title('成交量')
        ax2.set_ylabel('成交量')
        ax2.grid(True, alpha=0.3)
        
        # 3. MACD
        if 'macd' in data.columns:
            ax3 = axes[2]
            ax3.plot(data.index, data['macd'], label='MACD', color=self.colors[0])
            if 'macd_signal' in data.columns:
                ax3.plot(data.index, data['macd_signal'], label='Signal', color=self.colors[1])
            if 'macd_histogram' in data.columns:
                ax3.bar(data.index, data['macd_histogram'], alpha=0.6, 
                       color=['g' if x >= 0 else 'r' for x in data['macd_histogram']])
            ax3.set_title('MACD')
            ax3.set_ylabel('MACD')
            ax3.legend()
            ax3.grid(True, alpha=0.3)
        
        # 4. RSI
        if 'rsi' in data.columns:
            ax4 = axes[3]
            ax4.plot(data.index, data['rsi'], label='RSI', color=self.colors[4])
            ax4.axhline(y=70, color='r', linestyle='--', alpha=0.7, label='超买线(70)')
            ax4.axhline(y=30, color='g', linestyle='--', alpha=0.7, label='超卖线(30)')
            ax4.set_title('RSI')
            ax4.set_ylabel('RSI')
            ax4.set_xlabel('日期')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
            ax4.set_ylim(0, 100)
        
        plt.tight_layout()
        
        if save:
            filename = f"price_indicators_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()
    
    def plot_trading_signals(self, data: pd.DataFrame, signals: pd.Series,
                           title: str = "交易信号", save: bool = True) -> None:
        """绘制交易信号"""
        fig, ax = plt.subplots(figsize=(15, 8))
        
        # 绘制价格
        ax.plot(data.index, data['close'], label='收盘价', linewidth=1.5, color='blue')
        
        # 绘制买入信号
        buy_signals = signals == 1
        ax.scatter(data.index[buy_signals], data['close'][buy_signals], 
                  color='green', marker='^', s=100, label='买入信号', alpha=0.8)
        
        # 绘制卖出信号
        sell_signals = signals == -1
        ax.scatter(data.index[sell_signals], data['close'][sell_signals], 
                  color='red', marker='v', s=100, label='卖出信号', alpha=0.8)
        
        ax.set_title(title)
        ax.set_ylabel('价格')
        ax.set_xlabel('日期')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save:
            filename = f"trading_signals_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()
    
    def plot_backtest_results(self, backtest_data: Dict, title: str = "回测结果",
                            save: bool = True) -> None:
        """绘制回测结果"""
        fig, axes = plt.subplots(3, 1, figsize=(15, 12))
        
        dates = backtest_data['dates']
        prices = backtest_data['prices']
        balances = backtest_data['balances']
        positions = backtest_data['positions']
        
        # 1. 价格和仓位
        ax1 = axes[0]
        ax1_twin = ax1.twinx()
        
        ax1.plot(dates, prices, 'b-', label='价格', linewidth=1.5)
        ax1_twin.plot(dates, positions, 'r-', label='仓位', linewidth=2, alpha=0.7)
        
        ax1.set_ylabel('价格', color='b')
        ax1_twin.set_ylabel('仓位', color='r')
        ax1.set_title(f'{title} - 价格与仓位')
        ax1.grid(True, alpha=0.3)
        ax1.legend(loc='upper left')
        ax1_twin.legend(loc='upper right')
        
        # 2. 账户余额
        ax2 = axes[1]
        ax2.plot(dates, balances, 'g-', linewidth=2, label='账户余额')
        initial_balance = balances[0] if balances else 100000
        ax2.axhline(y=initial_balance, color='k', linestyle='--', alpha=0.5, label='初始资金')
        ax2.set_ylabel('账户余额 ($)')
        ax2.set_title('账户余额变化')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        ax2.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'${x:,.0f}'))
        
        # 3. 收益率
        if len(balances) > 1:
            returns = pd.Series(balances).pct_change().dropna()
            ax3 = axes[2]
            ax3.plot(dates[1:], returns.cumsum(), 'purple', linewidth=2, label='累计收益率')
            ax3.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            ax3.set_ylabel('累计收益率')
            ax3.set_xlabel('日期')
            ax3.set_title('累计收益率变化')
            ax3.grid(True, alpha=0.3)
            ax3.legend()
            ax3.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x:.1%}'))
        
        plt.tight_layout()
        
        if save:
            filename = f"backtest_results_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()
    
    def plot_algorithm_comparison(self, results: Dict[str, Dict], 
                                title: str = "算法性能对比", save: bool = True) -> None:
        """绘制算法对比图表"""
        if len(results) < 2:
            print("⚠️ 需要至少2个算法结果进行对比")
            return
        
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
        
        algos = list(results.keys())
        metrics = ['avg_return', 'avg_sharpe_ratio', 'avg_max_drawdown', 
                  'win_rate', 'avg_trades', 'std_return']
        metric_names = ['平均收益率', '夏普比率', '最大回撤', '胜率', '平均交易次数', '收益波动率']
        
        for i, (metric, name) in enumerate(zip(metrics, metric_names)):
            ax = axes[i//3, i%3]
            values = [results[algo][metric] for algo in algos]
            
            bars = ax.bar(algos, values, alpha=0.7, 
                         color=[self.colors[i % len(self.colors)] for i in range(len(algos))])
            
            # 添加数值标签
            for bar, value in zip(bars, values):
                height = bar.get_height()
                if metric in ['avg_return', 'avg_max_drawdown', 'win_rate', 'std_return']:
                    label = f'{value:.2%}'
                else:
                    label = f'{value:.3f}' if metric == 'avg_sharpe_ratio' else f'{value:.1f}'
                
                ax.text(bar.get_x() + bar.get_width()/2., height,
                       label, ha='center', va='bottom', fontweight='bold')
            
            ax.set_title(name, fontsize=12, fontweight='bold')
            ax.set_ylabel(name)
            
            # 格式化y轴
            if metric in ['avg_return', 'avg_max_drawdown', 'win_rate', 'std_return']:
                ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x:.1%}'))
            
            ax.grid(True, alpha=0.3)
        
        plt.suptitle(title, fontsize=16, fontweight='bold')
        plt.tight_layout()
        
        if save:
            filename = f"algorithm_comparison_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()
    
    def create_interactive_chart(self, data: pd.DataFrame, title: str = "交互式图表") -> None:
        """创建交互式图表"""
        try:
            # 创建子图
            fig = make_subplots(
                rows=3, cols=1,
                shared_xaxes=True,
                vertical_spacing=0.05,
                subplot_titles=('价格', '成交量', 'RSI'),
                row_width=[0.2, 0.1, 0.1]
            )
            
            # 价格图表
            fig.add_trace(
                go.Candlestick(
                    x=data.index,
                    open=data['open'],
                    high=data['high'],
                    low=data['low'],
                    close=data['close'],
                    name='价格'
                ),
                row=1, col=1
            )
            
            # 移动平均线
            if 'sma_20' in data.columns:
                fig.add_trace(
                    go.Scatter(
                        x=data.index,
                        y=data['sma_20'],
                        mode='lines',
                        name='SMA(20)',
                        line=dict(color='orange', width=1)
                    ),
                    row=1, col=1
                )
            
            # 成交量
            fig.add_trace(
                go.Bar(
                    x=data.index,
                    y=data['volume'],
                    name='成交量',
                    marker_color='lightblue'
                ),
                row=2, col=1
            )
            
            # RSI指标
            if 'rsi' in data.columns:
                fig.add_trace(
                    go.Scatter(
                        x=data.index,
                        y=data['rsi'],
                        mode='lines',
                        name='RSI',
                        line=dict(color='purple', width=1)
                    ),
                    row=3, col=1
                )
                
                # 添加超买超卖线
                fig.add_trace(
                    go.Scatter(
                        x=data.index,
                        y=[70] * len(data),
                        mode='lines',
                        name='超买线',
                        line=dict(color='red', width=1, dash='dash')
                    ),
                    row=3, col=1
                )
                
                fig.add_trace(
                    go.Scatter(
                        x=data.index,
                        y=[30] * len(data),
                        mode='lines',
                        name='超卖线',
                        line=dict(color='green', width=1, dash='dash')
                    ),
                    row=3, col=1
                )
            
            # 更新布局
            fig.update_layout(
                title=title,
                height=800,
                width=1200,
                showlegend=True,
                legend=dict(
                    orientation="h",
                    yanchor="bottom",
                    y=1.02,
                    xanchor="right",
                    x=1
                )
            )
            
            # 更新x轴
            fig.update_xaxes(
                rangeslider_visible=False,
                rangebreaks=[
                    dict(bounds=["sat", "mon"]),  # 隐藏周末
                ],
                row=3, col=1
            )
            
            # 显示图表
            fig.show()
            
            # 保存为HTML
            filename = f"interactive_{title.replace(' ', '_')}.html"
            fig.write_html(self.save_dir / filename)
            print(f"📈 交互式图表已保存: {filename}")
            
        except Exception as e:
            print(f"❌ 创建交互式图表失败: {e}")
    
    def plot_reward_history(self, rewards: List[float], 
                          title: str = "强化学习奖励历史", save: bool = True) -> None:
        """绘制奖励历史图表"""
        fig, ax = plt.subplots(figsize=(12, 6))
        
        episodes = list(range(1, len(rewards) + 1))
        
        # 绘制奖励
        ax.plot(episodes, rewards, 'b-', alpha=0.5, label='奖励')
        
        # 添加移动平均线
        window = min(100, max(5, len(rewards) // 10))
        rewards_series = pd.Series(rewards)
        rewards_ma = rewards_series.rolling(window=window).mean()
        ax.plot(episodes, rewards_ma, 'r-', linewidth=2, label=f'移动平均 (窗口={window})')
        
        # 添加趋势线
        try:
            from scipy import stats
            slope, intercept, r_value, p_value, std_err = stats.linregress(episodes, rewards)
            trend_line = [slope * x + intercept for x in episodes]
            ax.plot(episodes, trend_line, 'g--', linewidth=1, 
                   label=f'趋势线 (斜率={slope:.4f})')
        except:
            pass
        
        ax.set_title(title)
        ax.set_xlabel('回合数')
        ax.set_ylabel('奖励')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        # 添加统计信息
        text_info = (
            f"总回合数: {len(rewards)}\n"
            f"平均奖励: {np.mean(rewards):.2f}\n"
            f"最大奖励: {np.max(rewards):.2f}\n"
            f"最小奖励: {np.min(rewards):.2f}\n"
            f"标准差: {np.std(rewards):.2f}"
        )
        props = dict(boxstyle='round', facecolor='white', alpha=0.7)
        ax.text(0.02, 0.97, text_info, transform=ax.transAxes, fontsize=10,
               verticalalignment='top', bbox=props)
        
        plt.tight_layout()
        
        if save:
            filename = f"reward_history_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()
    
    def plot_learning_curves(self, training_data: Dict[str, List[float]], 
                           title: str = "学习曲线", save: bool = True) -> None:
        """绘制学习曲线"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        # 1. 奖励曲线
        if 'rewards' in training_data:
            ax1 = axes[0, 0]
            rewards = training_data['rewards']
            episodes = list(range(1, len(rewards) + 1))
            
            ax1.plot(episodes, rewards, 'b-', alpha=0.5)
            
            # 添加移动平均线
            window = min(100, max(5, len(rewards) // 10))
            rewards_series = pd.Series(rewards)
            rewards_ma = rewards_series.rolling(window=window).mean()
            ax1.plot(episodes, rewards_ma, 'r-', linewidth=2, 
                    label=f'移动平均 (窗口={window})')
            
            ax1.set_title('奖励曲线')
            ax1.set_xlabel('回合数')
            ax1.set_ylabel('奖励')
            ax1.legend()
            ax1.grid(True, alpha=0.3)
        
        # 2. 损失曲线
        if 'losses' in training_data:
            ax2 = axes[0, 1]
            losses = training_data['losses']
            steps = list(range(1, len(losses) + 1))
            
            ax2.plot(steps, losses, 'g-', alpha=0.5)
            
            # 添加移动平均线
            window = min(100, max(5, len(losses) // 10))
            losses_series = pd.Series(losses)
            losses_ma = losses_series.rolling(window=window).mean()
            ax2.plot(steps, losses_ma, 'r-', linewidth=2, 
                    label=f'移动平均 (窗口={window})')
            
            ax2.set_title('损失曲线')
            ax2.set_xlabel('训练步数')
            ax2.set_ylabel('损失')
            ax2.legend()
            ax2.grid(True, alpha=0.3)
        
        # 3. 探索率曲线
        if 'epsilons' in training_data:
            ax3 = axes[1, 0]
            epsilons = training_data['epsilons']
            episodes = list(range(1, len(epsilons) + 1))
            
            ax3.plot(episodes, epsilons, 'purple', linewidth=2)
            ax3.set_title('探索率 (Epsilon) 曲线')
            ax3.set_xlabel('回合数')
            ax3.set_ylabel('探索率')
            ax3.grid(True, alpha=0.3)
        
        # 4. 账户余额曲线
        # if 'balances' in training_1, 1]
        if 'balances' in training_1:
            balances = training_data['balances']
            episodes = list(range(1, len(balances) + 1))
            
            ax4.plot(episodes, balances, 'orange', linewidth=2)
            
            # 添加初始余额基准线
            initial_balance = balances[0] if balances else 100000
            ax4.axhline(y=initial_balance, color='k', linestyle='--', alpha=0.5, 
                       label='初始余额')
            
            ax4.set_title('账户余额曲线')
            ax4.set_xlabel('回合数')
            ax4.set_ylabel('账户余额')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
            ax4.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'${x:,.0f}'))
        
        plt.suptitle(title, fontsize=16, fontweight='bold')
        plt.tight_layout()
        
        if save:
            filename = f"learning_curves_{title.replace(' ', '_')}.png"
            plt.savefig(self.save_dir / filename, dpi=300, bbox_inches='tight')
            print(f"📈 图表已保存: {filename}")
        
        plt.show()

if __name__ == "__main__":
    # 测试可视化工具
    print("🧪 测试可视化工具")
    print("=" * 40)
    
    # 创建可视化器
    visualizer = TradingVisualizer()
    
    # 创建示例数据
    import sys
    from pathlib import Path
    
    # 添加项目路径
    project_root = Path(__file__).parent.parent
    sys.path.insert(0, str(project_root))
    
    from utils.data_loader import create_sample_data
    
    # 创建样本数据
    sample_data = create_sample_data("TEST", 100)
    
    # 添加技术指标
    from utils.data_loader import DataLoader
    
    loader = DataLoader()
    enhanced_data = loader.add_technical_indicators(sample_data)
    
    # 测试价格和指标图表
    visualizer.plot_price_and_indicators(enhanced_data, "示例数据")
    
    # 创建模拟交易信号
    signals = pd.Series(0, index=enhanced_data.index)
    
    # 简单的金叉死叉策略
    for i in range(5, len(enhanced_data)):
        if (enhanced_data['sma_5'].iloc[i-1] <= enhanced_data['sma_20'].iloc[i-1] and 
            enhanced_data['sma_5'].iloc[i] > enhanced_data['sma_20'].iloc[i]):
            signals.iloc[i] = 1  # 金叉买入
        elif (enhanced_data['sma_5'].iloc[i-1] >= enhanced_data['sma_20'].iloc[i-1] and 
              enhanced_data['sma_5'].iloc[i] < enhanced_data['sma_20'].iloc[i]):
            signals.iloc[i] = -1  # 死叉卖出
    
    # 测试交易信号图表
    visualizer.plot_trading_signals(enhanced_data, signals, "金叉死叉信号")
    
    # 创建模拟回测结果
    backtest_data = {
        'dates': enhanced_data.index,
        'prices': enhanced_data['close'].values,
        'balances': [100000 + i * 1000 for i in range(len(enhanced_data))],
        'positions': np.cumsum(signals.values)
    }
    
    # 测试回测结果图表
    visualizer.plot_backtest_results(backtest_data, "样本回测")
    
    # 测试奖励历史图表
    rewards = np.random.normal(50, 20, 200).cumsum()
    visualizer.plot_reward_history(rewards, "样本奖励")
    
    print("✅ 可视化工具测试完成")